# Copyright 2021-2023 @ Shenzhen Bay Laboratory &
#                       Peking University &
#                       Huawei Technologies Co., Ltd
#
# This code is a part of MindSPONGE:
# MindSpore Simulation Package tOwards Next Generation molecular modelling.
#
# MindSPONGE is open-source software based on the AI-framework:
# MindSpore (https://www.mindspore.cn/)
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ============================================================================
"""
Leap-frog integrator
"""

from typing import Tuple

import torch
import math
from torch import Tensor

from .integrator import Integrator, _integrator_register
from ..thermostat import Thermostat
from ..barostat import Barostat
from ..constraint import Constraint
from ...system import Molecule
from ...function import get_arguments


@_integrator_register('nmpimd')
class NMPIMD(Integrator):
    r"""A leap-frog integrator based on "middle scheme" developed by Jian Liu, et al. It is a subclass of `Integrator`.

    Reference:

        Zhang, Z.; Yan, K; Liu, X.; Liu, J.
        A Leap-Frog Algorithm-based Efficient Unified Thermostat Scheme for Molecular Dynamics [J].
        Chinese Science Bulletin, 2018, 63(33): 3467-3483.

    Args:

        system (Molecule):          Simulation system

        thermostat (Thermostat):    Thermostat for temperature coupling. Default: None

        barostat (Barostat):        Barostat for pressure coupling. Default: None

        constraint (Constraint):    Constraint algorithm. Default: None


    Supported Platforms:

        ``Ascend`` ``GPU``

    """
    def __init__(self,
                 system: Molecule,
                 time_step = 1e-3,
                 temperature = 300.,
                 gamma0 = 10.,
                 remove_com = True,
                 thermostat: Thermostat = None,
                 barostat: Barostat = None,
                 constraint: Constraint = None,
                 **kwargs
                 ):

        super().__init__(
            system=system,
            time_step=time_step,
            thermostat=thermostat,
            barostat=barostat,
            constraint=constraint,
        )
        self._kwargs = get_arguments(locals(), kwargs)
        self.name = 'PILE'
        n = self.num_walker
        self.temperature = self._get_mw_tensor(temperature, 'temperature')
        self.time_step = time_step
        boltzmann = system.units.boltzmann
        hbar = 1.054572e-37 * 6.02e23 * 1e12
        self.beta = 1. / (boltzmann * self.temperature)
        self.omega_p = math.sqrt(n) / (self.beta * hbar)
        self.gamma = self.omega_p*torch.ones(n,1,1,device=self.device)
        self.gamma0 = gamma0
        self.gamma[0] = self.gamma0
        self.c1 = torch.exp(-self.gamma*self.time_step)
        self.c2 = torch.sqrt(1-self.c1**2) / self.kinetic_unit_scale
        self.m_bar = torch.zeros_like(self.atom_mass,device=self.device).unsqueeze(-1).repeat(n,1,1)
        for i in range(1,n):
            self.m_bar[i] = (i+1)/i*self.atom_mass.unsqueeze(-1)
        self.m_tilde = torch.cat([self.atom_mass.unsqueeze(-1),self.m_bar[1:]])
        self.i = torch.arange(n,device=self.device).reshape(-1,1,1)
        self.i1 = self.i+1
        self.sqrt_m_tilde = torch.sqrt(self.m_tilde)
        self.remove_com = remove_com

    def forward(self,
                  coordinate: Tensor,
                  velocity: Tensor,
                  force: Tensor,
                  energy: Tensor,
                  kinetics: Tensor,
                  virial: Tensor = None,
                  pbc_box: Tensor = None,
                  step: int = 0,
                  ) -> Tuple[Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor]:

        # (B,A,D) = (B,A,D) * (B,A,1)
        x_tilde = coordinate - torch.roll(coordinate,-1,dims=0) * (self.i/self.i1) - coordinate[0]/self.i1
        x_tilde[0] = coordinate[0]
        v_tilde = velocity - torch.roll(velocity,-1,dims=0) * (self.i/self.i1) - velocity[0]/self.i1
        v_tilde[0] = velocity[0]
        f_tilde = 1/self.num_walker*torch.cumsum(self.i*force,0) / self.i
        f_tilde[0] = force.mean(0)

        p_tilde = v_tilde * self.m_tilde
        p_tilde_half = p_tilde + self.acc_unit_scale * f_tilde * self.time_step - self.omega_p ** 2 * x_tilde * self.time_step * self.m_bar
        # (B,A,D) = (B,A,D) - (B,1,D)
        if self.remove_com:
            p_tilde_half = p_tilde_half - self.get_com_velocity(p_tilde_half)
        kinetics = self.get_pimd_kinetics(p_tilde_half)
       
        # R(t+0.5) = R(t) + v(t+0.5) * dt
        x_tilde_half = x_tilde + p_tilde_half / self.m_tilde * 0.5 * self.time_step

        if self.temperature is not None:
            # v'(t+0.5) = f_T[v(t+0.5)]
            p_tilde_half = self.langevin_thermostat(p_tilde_half, step)

        # R(t+1) = R(t+0.5) + v'(t+0.5) * dt
        x_tilde_new = x_tilde_half + p_tilde_half / self.m_tilde * 0.5 * self.time_step
        v_tilde_new = p_tilde_half / self.m_tilde

        summation_x = x_tilde_new/self.i
        summation_x[0] = 0
        coordinate_new = x_tilde_new[0] + torch.flip(torch.cumsum(torch.flip(summation_x,dims=[0]),0),dims=[0])*self.i

        summation_v = v_tilde_new/self.i
        summation_v[0] = 0
        velocity_new = v_tilde_new[0] + torch.flip(torch.cumsum(torch.flip(summation_v,dims=[0]),0),dims=[0])*self.i

        return coordinate_new, velocity_new, force, energy, kinetics, virial, pbc_box
    
    def langevin_thermostat(self, p, step):
        if self.control_step == 1 or step % self.control_step == 0:
            p = self.c1 * p + self.c2 * torch.sqrt(1 / self.beta) * self.sqrt_m_tilde * torch.randn(p.shape, device=self.device)

        return p
    
    def get_pimd_kinetics(self, p):
        ke = p ** 2 / (2*self.m_tilde)
        kinetics = torch.sum(ke, -2)
        return kinetics
    
    def get_pimd_com_velocity(self, momentum: Tensor):
        com_vel = momentum.sum(-2,keepdim=True)/self.m_tilde.sum(-2,keepdim=True) * self.m_tilde
        return com_vel